{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TVM 模块插件\n",
    "\n",
    "```{topic} 导航\n",
    "建议暴露能够修改由调用者传递的内存并调用运行时 API 进行内存分配的函数，允许插件模块创建并返回托管对象。然而，需要记住以下几个限制：\n",
    "\n",
    "- 如果模块返回对象，需要确保在卸载模块之前销毁该对象。否则，由于调用未加载的析构函数，可能会导致段错误（segfault）。\n",
    "- 如果模块返回 PackedFunc，那么需要确保 DLL 的 libc 和 TVM 运行时库匹配。否则，由于 `std::function` 的不兼容性，可能会导致段错误（segfault）。\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "::::{dropdown} C++ 源码：\n",
    "```{literalinclude} src/testing/plugin_module.cc\n",
    ":language: C++\n",
    "```\n",
    "::::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "编译："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "g++ -std=c++17 -O2 -fPIC -I/media/pc/data/lxw/ai/tvm/include -I/media/pc/data/lxw/ai/tvm/3rdparty/dmlc-core/include -I/media/pc/data/lxw/ai/tvm/3rdparty/dlpack/include -Iinclude -DDMLC_USE_LOGGING_LIBRARY=\\<tvm/runtime/logging.h\\> -shared -o outputs/libs/libtvm_plugin_module.so src/testing/plugin_module.cc -ldl -pthread -L/media/pc/data/lxw/ai/tvm/build\n"
     ]
    }
   ],
   "source": [
    "!make outputs/libs/libtvm_plugin_module.so"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Python 端使用："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "mod = tvm.runtime.load_module(\"outputs/libs/libtvm_plugin_module.so\")\n",
    "\n",
    "# 普通函数\n",
    "assert mod[\"AddOne\"](10) == 11\n",
    "assert mod[\"SubOne\"](10) == 9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "高级用法——返回模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "mymod = mod[\"CreateMyModule\"](10)\n",
    "fadd = mymod[\"add\"]\n",
    "assert fadd(10) == 20\n",
    "assert mymod[\"mul\"](10) == 100"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
